import pandas as pd
import numpy as np
import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_selection import VarianceThreshold, mutual_info_classif
from sklearn.preprocessing import StandardScaler

# 下载NLTK依赖资源（首次运行需执行）
nltk.download('punkt')
nltk.download('stopwords')


class DataProcessor:
    def __init__(self):
        """
        初始化数据处理器
        工单编号：大数据-用户画像-18-京东评论情感预测
        功能：处理需求文档中的4类数据集（商品信息.csv、商品类别列表.csv、训练集、测试集），完成清洗、特征工程
        """
        # 中文停用词表
        self.stop_words = set(stopwords.words('chinese'))
        # TF-IDF向量器（文本特征衍生）
        self.tfidf = TfidfVectorizer(max_features=10000, tokenizer=word_tokenize)
        # 特征选择器
        self.var_selector = VarianceThreshold(threshold=0.1)  # 方差筛选
        self.scaler = StandardScaler()  # 数值特征标准化

    def load_data(self, data_paths):
        """
        加载数据集
        :param data_paths: 字典，包含各文件路径，格式：{'product_info': 'xxx.csv', 'category': 'xxx.csv', 'train': 'xxx.csv', 'test': 'xxx.csv'}
        :return: 字典，包含加载后的DataFrame
        工单编号：大数据-用户画像-18-京东评论情感预测
        """
        data = {}
        # 加载商品信息（需求文档4.2.1）
        data['product_info'] = pd.read_csv(data_paths['product_info'],
                                           usecols=['产品ID（PRODUCT_0）', '商品名', '所属类别'],
                                           dtype={'产品ID（PRODUCT_0）': str, '商品名': str, '所属类别': str})
        # 加载商品类别列表（需求文档4.2.2）
        data['category'] = pd.read_csv(data_paths['category'],
                                       usecols=['类别ID', '类别名称'],
                                       dtype={'类别ID': str, '类别名称': str})
        # 加载训练集（需求文档4.2.3）
        data['train'] = pd.read_csv(data_paths['train'],
                                    usecols=['数据ID', '用户ID', '商品ID', '评论时间戳', '评论标题', '评论内容',
                                             '评分'],
                                    dtype={'数据ID': str, '用户ID': int, '商品ID': str, '评论时间戳': int,
                                           '评论标题': str, '评论内容': str, '评分': int})
        # 加载测试集（需求文档4.2.4）
        data['test'] = pd.read_csv(data_paths['test'],
                                   usecols=['数据ID', '用户ID', '商品ID', '评论时间戳', '评论标题', '评论内容'],
                                   dtype={'数据ID': str, '用户ID': int, '商品ID': str, '评论时间戳': int,
                                          '评论标题': str, '评论内容': str})
        print("数据集加载完成")
        return data

    def clean_data(self, data):
        """
        数据清洗：处理缺失值、异常值、重复值（需求文档3.1数据处理功能）
        :param data: 字典，包含train和test的DataFrame
        :return: 字典，清洗后的DataFrame
        工单编号：大数据-用户画像-18-京东评论情感预测
        """
        # 1. 处理训练集
        train = data['train'].copy()
        # 缺失值处理：评论标题/内容填充“无评论”，其他关键字段缺失删除
        train['评论标题'] = train['评论标题'].fillna('无评论')
        train['评论内容'] = train['评论内容'].fillna('无评论')
        train = train.dropna(subset=['用户ID', '商品ID', '评论时间戳', '评分'])

        # 异常值处理：评分仅保留1-5分，时间戳限定2011-2014年（2011-01-01：1293840000；2014-03-31：1396272000）
        train = train[(train['评分'] >= 1) & (train['评分'] <= 5)]
        train = train[(train['评论时间戳'] >= 1293840000) & (train['评论时间戳'] <= 1396272000)]

        # 重复值处理：基于“数据ID+用户ID+商品ID”去重
        train = train.drop_duplicates(subset=['数据ID', '用户ID', '商品ID'])

        # 2. 处理测试集（逻辑与训练集一致，无评分字段）
        test = data['test'].copy()
        test['评论标题'] = test['评论标题'].fillna('无评论')
        test['评论内容'] = test['评论内容'].fillna('无评论')
        test = test.dropna(subset=['用户ID', '商品ID', '评论时间戳'])
        test = test[(test['评论时间戳'] >= 1293840000) & (test['评论时间戳'] <= 1396272000)]
        test = test.drop_duplicates(subset=['数据ID', '用户ID', '商品ID'])

        data['train_clean'] = train
        data['test_clean'] = test
        print("数据清洗完成，训练集样本数：{}，测试集样本数：{}".format(len(train), len(test)))
        return data

    def feature_engineering(self, data):
        """
        特征工程：文本特征衍生、数值特征衍生、特征选择（需求文档3.1数据处理功能）
        :param data: 字典，包含清洗后的train、test及商品信息
        :return: 字典，包含最终特征矩阵（X_train, X_test）和标签（y_train）
        工单编号：大数据-用户画像-18-京东评论情感预测
        """
        train = data['train_clean'].copy()
        test = data['test_clean'].copy()
        product_info = data['product_info'].copy()

        # 1. 文本特征衍生：合并评论标题+内容，生成TF-IDF特征
        train['文本合并'] = train['评论标题'] + ' ' + train['评论内容']
        test['文本合并'] = test['评论标题'] + ' ' + test['评论内容']

        # 训练TF-IDF并转换文本
        tfidf_train = self.tfidf.fit_transform(train['文本合并']).toarray()
        tfidf_test = self.tfidf.transform(test['文本合并']).toarray()
        tfidf_cols = ['tfidf_' + str(i) for i in range(tfidf_train.shape[1])]

        # 2. 数值特征衍生：时间特征（小时、星期、月份）、商品类目特征
        def extract_time_features(timestamp):
            """从时间戳提取时间特征"""
            dt = pd.to_datetime(timestamp, unit='s')
            return pd.Series([dt.hour, dt.dayofweek, dt.month], index=['小时', '星期', '月份'])

        # 训练集时间特征
        train_time = train['评论时间戳'].apply(extract_time_features)
        # 测试集时间特征
        test_time = test['评论时间戳'].apply(extract_time_features)

        # 商品类目特征：关联商品信息，提取一级类目（所属类别按“,”分割，取第一个）
        product_info['一级类目'] = product_info['所属类别'].str.split(',').str[0]
        train = train.merge(product_info[['产品ID（PRODUCT_0）', '一级类目']],
                            left_on='商品ID', right_on='产品ID（PRODUCT_0）', how='left')
        test = test.merge(product_info[['产品ID（PRODUCT_0）', '一级类目']],
                          left_on='商品ID', right_on='产品ID（PRODUCT_0）', how='left')

        # 一级类目独热编码
        train_category = pd.get_dummies(train['一级类目'], prefix='类目')
        test_category = pd.get_dummies(test['一级类目'], prefix='类目')
        # 确保测试集与训练集类目特征一致
        common_cols = set(train_category.columns) & set(test_category.columns)
        train_category = train_category[list(common_cols)]
        test_category = test_category[list(common_cols)]

        # 3. 合并所有特征
        # 训练集特征：TF-IDF + 时间特征 + 类目特征
        X_train = np.hstack([
            tfidf_train,
            train_time.values,
            train_category.values
        ])
        # 测试集特征
        X_test = np.hstack([
            tfidf_test,
            test_time.values,
            test_category.values
        ])
        # 标签（评分）
        y_train = train['评分'].values

        # 4. 特征选择：方差筛选 + 互信息筛选
        # 方差筛选（去除低方差特征）
        X_train_var = self.var_selector.fit_transform(X_train)
        X_test_var = self.var_selector.transform(X_test)
        print("方差筛选后特征数：{}".format(X_train_var.shape[1]))

        # 互信息筛选（保留与标签相关性高的特征，取前80%）
        mi = mutual_info_classif(X_train_var, y_train, random_state=42)
        mi_threshold = np.percentile(mi, 20)  # 保留前80%
        mi_mask = mi >= mi_threshold
        X_train_selected = X_train_var[:, mi_mask]
        X_test_selected = X_test_var[:, mi_mask]
        print("互信息筛选后特征数：{}".format(X_train_selected.shape[1]))

        # 5. 数值标准化
        X_train_scaled = self.scaler.fit_transform(X_train_selected)
        X_test_scaled = self.scaler.transform(X_test_selected)

        return {
            'X_train': X_train_scaled,
            'y_train': y_train,
            'X_test': X_test_scaled,
            'test_ids': test['数据ID'].values  # 测试集数据ID（用于后续输出预测结果）
        }


